package main

import (
	"context"
	"log"
	"strconv"
	"time"

	"github.com/mark3labs/mcp-go/mcp"
	"github.com/mark3labs/mcp-go/server"
)

func main() {
	// 创建 MCP Server
	s := server.NewMCPServer(
		"Streamable HTTP Demo",
		"1.0.0",
		server.WithToolCapabilities(false),
		server.WithRecovery(),
	)

	// 注册工具
	calculatorTool := mcp.NewTool("calculate",
		// 工具描述和参数定义
		mcp.WithDescription("Perform basic arithmetic operations"),
		mcp.WithString("operation",
			mcp.Required(),
			mcp.Description("The operation to perform (add, subtract, multiply, divide)"),
			mcp.Enum("add", "subtract", "multiply", "divide"),
		),
		mcp.WithNumber("x",
			mcp.Required(),
			mcp.Description("First number"),
		),
		mcp.WithNumber("y",
			mcp.Required(),
			mcp.Description("Second number"),
		),
	)

	// 添加处理逻辑
	s.AddTool(calculatorTool, func(ctx context.Context, request mcp.CallToolRequest) (*mcp.CallToolResult, error) {
		// 解析参数
		op, err := request.RequireString("operation")
		if err != nil {
			return mcp.NewToolResultError(err.Error()), nil
		}

		x, err := request.RequireFloat("x")
		if err != nil {
			return mcp.NewToolResultError(err.Error()), nil
		}

		y, err := request.RequireFloat("y")
		if err != nil {
			return mcp.NewToolResultError(err.Error()), nil
		}

		// 执行计算操作
		var result float64
		switch op {
		case "add":
			result = x + y
		case "subtract":
			result = x - y
		case "multiply":
			result = x * y
		case "divide":
			if y == 0 {
				return mcp.NewToolResultError("cannot divide by zero"), nil
			}
			result = x / y
		default:
			return mcp.NewToolResultError("unknown operation"), nil
		}

		// 返回结果
		return mcp.NewToolResultText(strconv.FormatFloat(result, 'f', 2, 64)), nil
	})

	httpServer := server.NewStreamableHTTPServer(s,
		server.WithHeartbeatInterval(30*time.Second),
	)

	addr := ":8080" // HTTP 服务监听的端口
	log.Printf("Starting Streamable HTTP MCP Server on %s", addr)
	// 启动 HTTP Server (streamableHttp 方式)
	if err := httpServer.Start(addr); err != nil {
		log.Fatalf("Failed to start HTTP server: %v", err)
	}
}
